import cplex
import readuflp
import sys
import os

def get_ind(i,j,N):
    return j*N + i + 1

#stub
def gen_uflp(filename):
    prob = cplex.Cplex()
    num_fac, num_cty, fx_cost, cn_cost = readuflp.read_uflp(filename)

    # add x_ij variables
    for j in range(num_cty):

        varnames = ["x" + str( get_ind(i,j, num_fac) ) for i in range(num_fac)]
        varcosts = cn_cost[j]
        vartypes = ["B"] * num_fac
        prob.variables.add(obj = varcosts,
                           types = vartypes,
                           names = varnames)
        
    # add facility variable y_i
    var_names = ["y" + str(i) for i in range(num_fac)]
    var_costs = fx_cost
    var_types = ["B"] * num_fac
    prob.variables.add(obj = var_costs,
                       types = var_types,
                       names = var_names)        

    # add assignment constraints
    for j in range(num_cty):
        var_names = ["x" + str( get_ind(i,j, num_fac) ) for i in range(num_fac)]
        cty_expr = [varnames, [1] * num_fac]
        prob.linear_constraints.add(lin_expr = [cty_expr],
                                    senses = ["E"],
                                    rhs = [1])


    
    # add associativity constraints
    for i in range(num_fac):
        var_names = ["y" + str(i)] + \
                    ["x" + str( get_ind(i,j, num_fac) ) for j in range(num_cty)]
        var_coefs = [ -1*num_cty ] + \
                    [1] * num_cty
        asso_expr = [var_names, var_coefs]
        prob.linear_constraints.add(lin_expr = [asso_expr],
                                    senses = ["L"],
                                    rhs = [0])        


    return prob

def batch_gen_uflp(dirname):
    for filename in os.listdir(dirname):
        prob = gen_uflp(dirname + filename)
        prob.write("uflp" + filename + ".lp")

batch_gen_uflp("./uniform/")

#prob = gen_uflp( sys.argv[1] )
#prob.write( sys.argv[1] + ".lp")
